import psycopg2
import pandas as pd
from datetime import datetime

# 数据库连接配置
DB_CONFIG = {
    'dbname': 'heytea_db',
    'user': 'heytea',  # 确保使用具有足够权限的用户
    'password': '666666',
    'host': '127.0.0.1',
    'port': '5432'
}

def fetch_data():
    conn = psycopg2.connect(**DB_CONFIG)
    cursor = conn.cursor()

    # 查询库存消耗数据
    inventory_query = """
    SELECT order_date, product_code, product_name, total_qty 
    FROM inventory_date_summary
    WHERE EXTRACT(MONTH FROM order_date) = 3;
    """
    cursor.execute(inventory_query)
    inventory_data = cursor.fetchall()
    inventory_df = pd.DataFrame(inventory_data, columns=['order_date', 'product_code', 'product_name', 'total_qty'])

    # 查询天气数据
    weather_query = """
    SELECT date, highest_temperature, lowest_temperature, average_temperature, body_temperature,
           average_humidity, weather, precipitation_24h
    FROM daily_weather
    WHERE EXTRACT(MONTH FROM date) = 3;
    """
    cursor.execute(weather_query)
    weather_data = cursor.fetchall()
    weather_df = pd.DataFrame(weather_data, columns=['date', 'highest_temperature', 'lowest_temperature', 
                                                     'average_temperature', 'body_temperature', 'average_humidity', 
                                                     'weather', 'precipitation_24h'])

    cursor.close()
    conn.close()

    return inventory_df, weather_df

inventory_df, weather_df = fetch_data()

# 合并数据集
merged_df = pd.merge(inventory_df, weather_df, left_on='order_date', right_on='date', how='inner')
merged_df.drop('date', axis=1, inplace=True)

# 添加星期几特征
merged_df['day_of_week'] = merged_df['order_date'].dt.dayofweek

# 将日期转换为时间戳
merged_df['order_date'] = merged_df['order_date'].apply(lambda x: x.timestamp())

# 对天气进行编码
merged_df['weather'] = merged_df['weather'].astype('category').cat.codes

# 查看数据
print(merged_df.head())

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

# 选择特征和目标变量
features = ['order_date', 'product_code', 'highest_temperature', 'lowest_temperature', 
            'average_temperature', 'body_temperature', 'average_humidity', 'weather', 
            'precipitation_24h', 'day_of_week']
target = 'total_qty'

X = merged_df[features]
y = merged_df[target]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练模型
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

# 评估模型
mse = mean_squared_error(y_test, y_pred)
print(f"Mean Squared Error: {mse}")